import torch
from torch import nn
from torch import fft
import math
from PIL import Image
from torchvision import transforms


def get_weight_matrix(H, W):
    x = torch.arange(H).unsqueeze(0)
    y = torch.arange(W).unsqueeze(0)
    u = (x * (H / (2 * torch.tensor(math.pi))))
    v = (y * (W / (2 * torch.tensor(math.pi)))).permute(1, 0)
    weight_matrix = torch.sqrt(u ** 2 + v ** 2) / (H * W)
    return weight_matrix.unsqueeze(0).unsqueeze(0)


class WFDLoss(nn.Module):

    def __init__(self, H, W):
        super(WFDLoss, self).__init__()
        self.register_buffer('weight_matrix', get_weight_matrix(H, W))

    def forward(self, origin, rec):
        """

        :param origin: 初始图像
        :param rec: 重构图像
        :return:
        """

        dft_origin = fft.fftn(origin, dim=(-2, -1))
        dft_rec = fft.fftn(rec, dim=(-2, -1))
        d_dft = torch.abs(dft_origin - dft_rec)
        d_weight_dft = self.weight_matrix * d_dft
        return d_weight_dft.mean()


if __name__ == '__main__':
    transform_ = transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
        transforms.Normalize([0.0], [0.5]),
    ])

    img = Image.open("./bottle.png").convert('L')
    img_tensor = transform_(img).unsqueeze(0)
    noise = torch.randn(img_tensor.size())
    std = 0.5
    img_tensor2 = img_tensor + std * noise

    loss = WFDLoss(256, 256)
    print(loss(img_tensor, img_tensor2))
